import torch
from torch.onnx import TrainingMode
from torch.utils.data import TensorDataset, DataLoader
from tvts.examples.vgg16_torch_cifar10_tvts import Vgg16
from PyCmpltrtok.common import *
from PyCmpltrtok.common_np import *
import os
import PyCmpltrtok.data.cifar10.load_cifar10 as cifar10
from tvts.examples.vgg16_torch_cifar10_tvts import compile, accuracy, evaluate


def _export():
    model_path = '/home/asuspei/PycharmProjects/GitHubAsusCondaP37tf24/tvts/examples/_save/vgg16_torch_cifar10_tvts.py/v2.0/vgg16_torch_cifar10-17-17.pth'
    VER = 'v1.0'

    # select device
    sep('cpu or gpu')
    device_id = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    print(device_id)
    device = torch.device(device_id)
    print(device)

    # the model
    sep('The model')
    model = Vgg16(10, (3, 32, 32)).to(device)
    sdict = torch.load(model_path)
    model.load_state_dict(sdict)

    # test model loaded
    N_BATCH_SIZE = 256
    _, _, x_test, y_test, label_names = cifar10.load()
    shape_ = cifar10.shape_
    x_test = x_test.reshape(-1, *shape_)
    x_test = uint8_to_flt_by_lut(x_test)
    x_test = torch.Tensor(x_test)
    y_test = torch.Tensor(y_test)
    ds_test = TensorDataset(x_test, y_test)
    dl_test = DataLoader(ds_test, N_BATCH_SIZE, drop_last=False)
    model_dict = compile(model=model, device=device, metric=accuracy)
    evaluate(model_dict, dl_test)

    # dir to save
    sep('The dir to save')
    BASE_DIR, FILE_NAME = os.path.split(__file__)
    PWD = os.getcwd()
    BASE_DIR = os.path.join(PWD, BASE_DIR)
    SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
    os.makedirs(SAVE_DIR, exist_ok=True)

    # export
    sep('Export')
    M = 64
    dummy_input = torch.randn(M, 3, 32, 32, device=device)
    tgt_path = os.path.join(SAVE_DIR, f'{rand_name_on_now()}-{M}.onnx')
    print(f'To: {tgt_path}')
    torch.onnx.export(
        model, dummy_input, tgt_path,
        verbose=True,
        training=TrainingMode.EVAL,
        input_names=['input'],
        output_names=['output'],
        keep_initializers_as_inputs=True,
    )


if '__main__' == __name__:
    _export()
